'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-09-23
Copyright © Department of Physics, Tsinghua University. All rights reserved

Define common types for AI tool
'''
ALL_CATEGORIES = []
with open("all-categories.txt", 'r') as fp:
    for line_id, line in enumerate(fp):
        ALL_CATEGORIES.append(line.replace('\n', ''))

ALL_CATEGORY_STR = ""
for line_id, line in enumerate(ALL_CATEGORIES):
    ALL_CATEGORY_STR += chr(ord('A') + line_id) + '.' + line + '\n'


def check_classification(response_str: str) -> tuple[bool, str, str]:
    ''' check the result of classification
        input:
            response_str: raw response string from LLM
        output:
            is_successful: bool, whether the classification is successful
            tag: tag of the class
            category_name: full name of the category
    '''
    response_str = response_str.replace(" ", "").replace("\n", "")
    if (
        len(response_str) > 1
        or (response_str[0] < 'A')
        or (ord(response_str[0]) > ord('A') + len(ALL_CATEGORIES) - 1)
    ):
        # failed
        return False, response_str, None

    else:
        tag = response_str
        category_name = (
            tag + '.' + ALL_CATEGORIES[ord(tag) - ord('A')])
        category_name = (category_name
                         .replace("\'", '')
                         .replace(" ", "_"))
        return True, tag, category_name


class AIBase:
    ''' Base class for all AI interfaces '''
    user_key: str
    model_name: str
    url: str

    def __init__(
            self,
            user_key: str,
            model_name: str,
            url: str
    ) -> None:
        self.user_key = user_key
        self.model_name = model_name
        self.url = url

    def get_content(self, messages: list[dict[str, str]]) -> str:
        ''' Get LLM response from messages, override
            messages: messages of LLM api
        '''
        pass

    def batch_run(self, message_list: list[list[dict]],
                  source_file: str, target_file: str):
        ''' Get LLM response by batch API, override
            message_list: list of messages
        '''
        pass

    def get_literature_classification_messages(self, literature_dict: dict) -> list[dict, dict]:
        ''' Get the message for literature classification '''
        system_prompt = (
            "[Character settings]You are a physicist and English native speaker."
            " Now you need to read the title and the abstract of an article, and"
            " tell me what category the article belongs to.\n\n\n "
            "[order]1.In every quiry you must return the letter of the most suitable "
            "category among the following choices: \n" + ALL_CATEGORY_STR + "\n "
            "2.Return the label of each choice only. **DO NOT RETURN FULL CAPTION**\n "
            "3.Each query is given in the format '[Title]<article-title>, "
            "[Abstract]<article-abstract>'\n\n\n"
            " [Example]A\nB\nC\nD\n\n"
        )
        content = "[Title]" + literature_dict["title"] + ", [Abstract]" + literature_dict["abstract"]
        return [{
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": content
                }]

    def classify_literature(self, literature_dict: dict) -> str:
        ''' Classify dict by title and abstract '''
        return self.get_content(
            self.get_literature_classification_messages(literature_dict)
        )

    def batch_classify(self, literature_dict_list: dict) -> list[str]:
        ''' Batch classify the literatures '''
        message_list = [
            self.get_literature_classification_messages(item)
            for item in literature_dict_list
        ]
        return self.batch_run(message_list, "data/temp.jsonl", "data/result.jsonl")


def main_classify_literatures(
    ai_model: AIBase,
    item_list: list[dict],
    max_trial: int = 3
) -> tuple[list, list]:
    ''' main function for classify the literatures
        input:
            ai_model: LLM model to classify
            item_list: list of literature dict
            max_trial: maximum trials for failed items
        output:
            success_list: list of succeed items. Format:
                (item, tag, category_name)
            failed_list: list of failed items. Format:
                (item, message)
    '''
    success_list = []
    print("Total items:", len(item_list))

    curr_items = item_list
    for run_turn_id in range(max_trial):
        print("Current trial:", run_turn_id)
        failed_list = []

        # check response
        for item_ind in range(len(curr_items)):
            response_str = ai_model.classify_literature(curr_items[item_ind])
            flag, tag, category_name = check_classification(response_str)
            if flag:    # successful
                print(" " * 100, end='\r')
                print("Item: %d, successful, %s" % (item_ind, category_name), end='\r')
                success_list.append(
                    (curr_items[item_ind], tag, category_name)
                )
            else:       # Failed
                print(" " * 100, end='\r')
                print("Item: %d, failed, %s" % (item_ind, response_str), end='\r')
                failed_list.append(
                    (curr_items[item_ind], response_str)
                )

        print(" " * 100, end='\r')
        print("Successful: %d, failed: %d" % (len(success_list), len(failed_list)))

        if not failed_list:
            break

        # update curr_items
        curr_items = [pair[0] for pair in failed_list]

    return success_list, failed_list
